package icu.nil.demo.shell.service

import com.fasterxml.jackson.databind.ObjectMapper
import com.fasterxml.jackson.module.kotlin.readValue
import com.jcraft.jsch.Channel
import com.jcraft.jsch.JSch
import com.jcraft.jsch.JSchException
import icu.nil.demo.shell.bean.SSHConnectInfo
import icu.nil.demo.shell.bean.WebSSHData
import icu.nil.demo.shell.config.USER_UUID_KEY
import icu.nil.demo.shell.config.WEBSSH_OPERATE_COMMAND
import icu.nil.demo.shell.config.WEBSSH_OPERATE_CONNECT
import icu.nil.demo.shell.config.logger
import org.springframework.stereotype.Service
import org.springframework.web.socket.TextMessage
import org.springframework.web.socket.WebSocketSession
import java.io.IOException
import java.util.concurrent.ConcurrentHashMap
import java.util.concurrent.Executors

/**
 * Create by xiangliyou on 2020/3/10 17:03.
 * email: xiangliyou@aliyun.com
 * Des:
 */
@Service
class WebSSHServiceImpl(
        private val mapper: ObjectMapper
) : WebSSHService {

    /**
     * 存放ssh连接信息的map
     */
    private val sshMap: MutableMap<String, SSHConnectInfo> = ConcurrentHashMap()

    /**
     * 线程池
     */
    private val executorService = Executors.newCachedThreadPool()

    /**
     * 初始化连接
     */
    override fun initConnection(session: WebSocketSession) {
        val sshConnectInfo = SSHConnectInfo(JSch(), session)

        val uuid = session.attributes[USER_UUID_KEY] as? String
                ?: throw RuntimeException("初始化连接失败")
        //将这个ssh连接信息放入map中
        sshMap[uuid] = sshConnectInfo
    }

    /**
     * 处理客户端发送的数据
     */
    override fun recvHandle(buffer: String, session: WebSocketSession) {
        val sshData: WebSSHData = mapper.readValue(buffer)
                ?: throw RuntimeException("Json转换异常")
        val userId = session.attributes[USER_UUID_KEY]
                ?: throw RuntimeException("userId not found")
        when (sshData.operate) {
            WEBSSH_OPERATE_CONNECT -> {
                val sshConnectInfo = sshMap[userId]
                        ?: throw RuntimeException("ConnectInfo not found")
                executorService.execute {
                    try {
                        connectToSSH(sshConnectInfo, sshData, session)
                    } catch (e: JSchException) {
                        close(session)
                        e.printStackTrace()
                    } catch (e: IOException) {
                        close(session)
                        e.printStackTrace()
                    }
                }
            }
            WEBSSH_OPERATE_COMMAND -> {
                val sshConnectInfo = sshMap[userId]
                        ?: throw RuntimeException("ConnectInfo not found")
                val channel = sshConnectInfo.channel
                if (channel != null) {
                    transToSSH(channel, sshData.command)
                } else {
                    close(session)
                }
            }
            else -> {
                logger().error("不支持的操作")
                close(session)
            }
        }
    }

    /**
     * 发送数据到客户端
     */
    override fun sendMessage(session: WebSocketSession, buffer: ByteArray) {
        session.sendMessage(TextMessage(buffer))
    }

    /**
     * 关闭连接
     */
    override fun close(session: WebSocketSession) {
        val userId = session.attributes[USER_UUID_KEY] as? String ?: return
        val sshConnectInfo = sshMap[userId] ?: return
        // 断开 ssh 连接
        sshConnectInfo.channel?.disconnect()
        sshMap.remove(userId)

        logger().error("exit-status: page connection close.")
    }

    /**
     * 连接到终端
     */
    private fun connectToSSH(sshConnectInfo: SSHConnectInfo, sshData: WebSSHData, session: WebSocketSession) {
        // 获得会话
        val jschSession = sshConnectInfo.jSch.getSession(sshData.username, sshData.host, sshData.port)
        jschSession.setConfig("StrictHostKeyChecking", "no")
        jschSession.setConfig("PreferredAuthentications", "keyboard-interactive,password,publickey")
        jschSession.setConfig("ConnectTime", "30000")

        jschSession.setPassword(sshData.password)
        // 连接, 30s 超时
        jschSession.connect()
        // 开启 shell 通道
        val channel = jschSession.openChannel("shell")
        channel.inputStream = null
        channel.connect()
        // 保存通道
        sshConnectInfo.channel = channel

        // 读取终端返回的信息
        try {
            val inputStream = channel.inputStream
            val buffer = ByteArray(1024)
            // 循环读取完返回的结果
            val result = mutableListOf<ByteArray>()
            var i: Int
            while (true) {
                Thread.sleep(10)
                while (inputStream.available() > 0) {
                    i = inputStream.read(buffer)
                    if (i < 0) {
                        break
                    }
                    // 获取有效数据
                    val readData = buffer.sliceArray(0 until i)
                    result.add(readData)
                }

                if (result.isNotEmpty()) {
                    // 合并读取到的数据并发送给前台
                    val data = result.reduce { acc, bytes -> acc.plus(bytes) }
                    sendMessage(session, data)

                    result.clear()
                }

                // 连接断开
                if (channel.isClosed) {
                    if (inputStream.available() > 0) {
                        continue
                    }

                    logger().error("exit-status: " + channel.exitStatus)
                    break
                }

            }
        } catch (e: Exception) {
            logger().error(e.localizedMessage, e)
        } finally {
            jschSession.disconnect()
            channel.disconnect()
        }
    }

    /**
     * 发送信息到终端
     */
    private fun transToSSH(channel: Channel, command: String) {
        val writer = channel.outputStream.bufferedWriter()

        writer.write(command)
        writer.flush()
    }
}